#include "HookLibrary.h"
#include "HookInfoFile.h"
#include "MessageWindow.h"

static volatile LPWSTR     HookInfoPath;

volatile PHOOK_INFO HookInfo;

LPWSTR WINAPI GetHookInfoPath(HWND hWndOwner)
{
	BOOL Succeeded = TRUE;
	LPWSTR Temp    = NULL;

	if (HookInfoPath == NULL)
	{
		HookInfoPath = BGHeapAllocStr(MAX_PATH);
		if (HookInfoPath == NULL)
		{
			Succeeded = FALSE;
			goto Exit;
		}

		if (!GetHookInfoDirectory(HookInfoPath))
		{
			Succeeded = FALSE;
			goto Exit;
		}

		Temp = BGHeapAllocStr(MAX_PATH);
		if (Temp == NULL)
		{
			Succeeded = FALSE;
			goto Exit;
		}

		if (GetHookInfoFileName(hWndOwner, Temp, MAX_PATH) == 0)
		{
			Succeeded = FALSE;
			goto Exit;
		}

		Succeeded = PathAppendW(HookInfoPath, Temp);
	}

Exit:

	if (!Succeeded)
	{
		if (HookInfoPath != NULL)
		{
			BGHeapFree(HookInfoPath);
			HookInfoPath = NULL;
		}
	}

	if (Temp != NULL)
	{
		BGHeapFree(Temp);
		Temp = NULL;
	}

	return HookInfoPath;
}

BOOL WINAPI GetHookInfoDirectory(LPWSTR lpPath)
{
	if (SUCCEEDED(SHGetFolderPathW(NULL, CSIDL_LOCAL_APPDATA, 
				    			   NULL, SHGFP_TYPE_CURRENT, 
								   lpPath)))
	{
		return PathAppendW(lpPath, L"BGShell");
	}

	return FALSE;
}

DWORD WINAPI GetHookInfoFileName(HWND hWnd, LPWSTR Path, DWORD cchPath)
{
	static LPCVOID FormatString = L"HookInfo-%1!u!-0x%2!08x!.bin";
	
	DWORD     ProcessId = 0;
	DWORD     SessionId = 0;
	DWORD_PTR Arguments[2];

	if (GetWindowThreadProcessId(hWnd, &ProcessId) != 0)
	{
		ProcessIdToSessionId(ProcessId, &SessionId);
	}

	Arguments[0] = (DWORD_PTR)(SessionId);
	Arguments[1] = (DWORD_PTR)(hWnd);

	return FormatMessageW(FORMAT_MESSAGE_FROM_STRING | FORMAT_MESSAGE_ARGUMENT_ARRAY,
		                  FormatString,
			        	  0, 0,
				          Path, cchPath,
					      (va_list*)(Arguments));
}

DWORD WINAPI ParseHookInfoKeys(CONST DWORD * First, DWORD ** Out)
{
	SIZE_T Size;
	SIZE_T Count;
	DWORD* Buffer;

	if (First == NULL)
	{
		return FALSE;
	}

	if (Out == NULL)
	{
		return FALSE;
	}

	Size   = 0;
	Count  = 0;
	Buffer = NULL;
	*Out   = NULL;

	while (*First != 0)
	{
		if (Count >= Size)
		{
			DWORD  NewSize = Size + 256;
			DWORD* NewBuffer = BGHeapRealloc(Buffer, (sizeof(DWORD) * NewSize));
			
			if (NewBuffer == NULL)
			{
				if (Buffer != NULL)
				{
					BGHeapFree(Buffer);
					Buffer = NULL;
				}

				return (DWORD)(-1);
			}
			
			Size   = NewSize;
			Buffer = NewBuffer;
		}

		Buffer[Count] = *First;

		++Count;
		++First;
	}

	*Out = Buffer;
	return Count;
}

BOOL WINAPI ParseHookInfo(LPCVOID MemorySection, PHOOK_INFO * Info)
{
	BOOL       Succeeded = TRUE;
	PHOOK_INFO TempInfo  = NULL;
	
	TempInfo = BGHeapAlloc(sizeof(HOOK_INFO));
	if (TempInfo == NULL)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	TempInfo->KeyCount = ParseHookInfoKeys((const DWORD *)(MemorySection), &(TempInfo->Keys));
	if (TempInfo->KeyCount == (DWORD)(-1))
	{
		Succeeded = FALSE;
		goto Exit;
	}

Exit:

	if (!Succeeded)
	{
		if (TempInfo != NULL)
		{
			if (TempInfo->Keys != NULL)
			{
				BGHeapFree(TempInfo->Keys);
				TempInfo->Keys = NULL;
			}

			BGHeapFree(TempInfo);
			TempInfo = NULL;
		}
	}

	*Info = TempInfo;
	return Succeeded;
}

BOOL WINAPI ReadHookInfo(LPCWSTR Path, PHOOK_INFO * Info)
{
	HANDLE File     = INVALID_HANDLE_VALUE;
	DWORD  FileSize = 0;

	LPCVOID View      = NULL;
	HANDLE  Mapping   = NULL;
	BOOL    Succeeded = TRUE;
	
	if (Info == NULL)
	{
		return FALSE;
	}

	*Info = NULL;

	File = CreateFileW(Path, 
		               GENERIC_READ,
				       FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
				       NULL,
				       OPEN_EXISTING,
              		   FILE_ATTRIBUTE_NORMAL,
    			       NULL);
	if (File == INVALID_HANDLE_VALUE)
	{
		Succeeded = FALSE;
		goto Exit;
	}


	FileSize = GetFileSize(File, NULL);
	if (FileSize == 0)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	Mapping = CreateFileMappingW(File,
		                         NULL,
								 PAGE_READONLY,
								 0, 0,
								 NULL);
	if (Mapping == NULL)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	View = MapViewOfFile(Mapping,
						 FILE_MAP_READ,
						 0, 0, 0);
	if (View == NULL)
	{
		Succeeded = FALSE;
		goto Exit;
	}

	Succeeded = ParseHookInfo(View, Info);

Exit:

	if (View != NULL)
	{
		Succeeded &= UnmapViewOfFile(View);
		View       = NULL;
	}

	if (Mapping != NULL)
	{
		Succeeded &= CloseHandle(Mapping);
		Mapping    = NULL;
	}

	if (File != INVALID_HANDLE_VALUE)
	{
		Succeeded &= CloseHandle(File);
		File       = INVALID_HANDLE_VALUE;
	}

	return Succeeded;
}
VOID WINAPI OnHookInfoChanged(HWND hWnd)
{
	EnterCriticalSection(&g_csModule);
	{
		LPCWSTR FilePath = GetHookInfoPath(hWnd);

		if (FilePath != NULL)
		{
			PHOOK_INFO OldHookInfo = HookInfo;
			PHOOK_INFO NewHookInfo = NULL;
			
			if (ReadHookInfo(FilePath, &NewHookInfo))
			{
				HookInfo = NewHookInfo;
				BGHeapFree(OldHookInfo);
			}
		}
	}
	LeaveCriticalSection(&g_csModule);
}